{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "f833faf5-ddab-4355-aa7f-010d8fc3fe84",
   "metadata": {},
   "source": [
    "Lost in the middle"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b5e6bca1-fe13-4aa5-b224-1d74feaf8482",
   "metadata": {},
   "outputs": [],
   "source": [
    "! pip install sentence-transformers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b4e1f7f2-865d-4188-9a3c-e54cb40848d3",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "C:\\Users\\Tingwang.Gao\\AppData\\Local\\Temp\\ipykernel_35288\\2321899963.py:9: LangChainDeprecationWarning: The class `HuggingFaceBgeEmbeddings` was deprecated in LangChain 0.2.2 and will be removed in 1.0. An updated version of the class exists in the :class:`~langchain-huggingface package and should be used instead. To use it run `pip install -U :class:`~langchain-huggingface` and import as `from :class:`~langchain_huggingface import HuggingFaceEmbeddings``.\n",
      "  embedings = HuggingFaceBgeEmbeddings(model_name=\"all-MiniLM-L6-v2\")\n",
      "D:\\py\\ai\\langchain-env\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "ename": "ImportError",
     "evalue": "Could not import chromadb python package. Please install it with `pip install chromadb`.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mModuleNotFoundError\u001b[0m                       Traceback (most recent call last)",
      "File \u001b[1;32mD:\\py\\ai\\langchain-env\\Lib\\site-packages\\langchain_community\\vectorstores\\chroma.py:83\u001b[0m, in \u001b[0;36mChroma.__init__\u001b[1;34m(self, collection_name, embedding_function, persist_directory, client_settings, collection_metadata, client, relevance_score_fn)\u001b[0m\n\u001b[0;32m     82\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m---> 83\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mchromadb\u001b[39;00m\n\u001b[0;32m     84\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mchromadb\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconfig\u001b[39;00m\n",
      "\u001b[1;31mModuleNotFoundError\u001b[0m: No module named 'chromadb'",
      "\nDuring handling of the above exception, another exception occurred:\n",
      "\u001b[1;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[2], line 24\u001b[0m\n\u001b[0;32m      9\u001b[0m embedings \u001b[38;5;241m=\u001b[39m HuggingFaceBgeEmbeddings(model_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mall-MiniLM-L6-v2\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     11\u001b[0m text \u001b[38;5;241m=\u001b[39m [\n\u001b[0;32m     12\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m篮球是一项伟大的运动。\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m     13\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m带我飞往月球是我最喜欢的歌曲之一。\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m     21\u001b[0m     \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m拉里.伯德是一位标志性的NBA球员\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m     22\u001b[0m ]\n\u001b[1;32m---> 24\u001b[0m retrieval \u001b[38;5;241m=\u001b[39m \u001b[43mChroma\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mfrom_texts\u001b[49m\u001b[43m(\u001b[49m\u001b[43mtext\u001b[49m\u001b[43m,\u001b[49m\u001b[43membedings\u001b[49m\u001b[43m)\u001b[49m\u001b[38;5;241m.\u001b[39mas_retriever(\n\u001b[0;32m     25\u001b[0m     search_kwargs\u001b[38;5;241m=\u001b[39m{\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mk\u001b[39m\u001b[38;5;124m\"\u001b[39m: \u001b[38;5;241m10\u001b[39m}\n\u001b[0;32m     26\u001b[0m )\n\u001b[0;32m     27\u001b[0m query \u001b[38;5;241m=\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m关于凯尔特人队你都知道什么?\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m     29\u001b[0m \u001b[38;5;66;03m#根据相关性返回文本块\u001b[39;00m\n",
      "File \u001b[1;32mD:\\py\\ai\\langchain-env\\Lib\\site-packages\\langchain_community\\vectorstores\\chroma.py:817\u001b[0m, in \u001b[0;36mChroma.from_texts\u001b[1;34m(cls, texts, embedding, metadatas, ids, collection_name, persist_directory, client_settings, client, collection_metadata, **kwargs)\u001b[0m\n\u001b[0;32m    784\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[0;32m    785\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mfrom_texts\u001b[39m(\n\u001b[0;32m    786\u001b[0m     \u001b[38;5;28mcls\u001b[39m: Type[Chroma],\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    796\u001b[0m     \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs: Any,\n\u001b[0;32m    797\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Chroma:\n\u001b[0;32m    798\u001b[0m \u001b[38;5;250m    \u001b[39m\u001b[38;5;124;03m\"\"\"Create a Chroma vectorstore from a raw documents.\u001b[39;00m\n\u001b[0;32m    799\u001b[0m \n\u001b[0;32m    800\u001b[0m \u001b[38;5;124;03m    If a persist_directory is specified, the collection will be persisted there.\u001b[39;00m\n\u001b[1;32m   (...)\u001b[0m\n\u001b[0;32m    815\u001b[0m \u001b[38;5;124;03m        Chroma: Chroma vectorstore.\u001b[39;00m\n\u001b[0;32m    816\u001b[0m \u001b[38;5;124;03m    \"\"\"\u001b[39;00m\n\u001b[1;32m--> 817\u001b[0m     chroma_collection \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;43mcls\u001b[39;49m\u001b[43m(\u001b[49m\n\u001b[0;32m    818\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcollection_name\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcollection_name\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    819\u001b[0m \u001b[43m        \u001b[49m\u001b[43membedding_function\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43membedding\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    820\u001b[0m \u001b[43m        \u001b[49m\u001b[43mpersist_directory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mpersist_directory\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    821\u001b[0m \u001b[43m        \u001b[49m\u001b[43mclient_settings\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclient_settings\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    822\u001b[0m \u001b[43m        \u001b[49m\u001b[43mclient\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mclient\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    823\u001b[0m \u001b[43m        \u001b[49m\u001b[43mcollection_metadata\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mcollection_metadata\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    824\u001b[0m \u001b[43m        \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m,\u001b[49m\n\u001b[0;32m    825\u001b[0m \u001b[43m    \u001b[49m\u001b[43m)\u001b[49m\n\u001b[0;32m    826\u001b[0m     \u001b[38;5;28;01mif\u001b[39;00m ids \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m    827\u001b[0m         ids \u001b[38;5;241m=\u001b[39m [\u001b[38;5;28mstr\u001b[39m(uuid\u001b[38;5;241m.\u001b[39muuid4()) \u001b[38;5;28;01mfor\u001b[39;00m _ \u001b[38;5;129;01min\u001b[39;00m texts]\n",
      "File \u001b[1;32mD:\\py\\ai\\langchain-env\\Lib\\site-packages\\langchain_core\\_api\\deprecation.py:216\u001b[0m, in \u001b[0;36mdeprecated.<locals>.deprecate.<locals>.finalize.<locals>.warn_if_direct_instance\u001b[1;34m(self, *args, **kwargs)\u001b[0m\n\u001b[0;32m    214\u001b[0m     warned \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[0;32m    215\u001b[0m     emit_warning()\n\u001b[1;32m--> 216\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mwrapped\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "File \u001b[1;32mD:\\py\\ai\\langchain-env\\Lib\\site-packages\\langchain_community\\vectorstores\\chroma.py:86\u001b[0m, in \u001b[0;36mChroma.__init__\u001b[1;34m(self, collection_name, embedding_function, persist_directory, client_settings, collection_metadata, client, relevance_score_fn)\u001b[0m\n\u001b[0;32m     84\u001b[0m     \u001b[38;5;28;01mimport\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01mchromadb\u001b[39;00m\u001b[38;5;21;01m.\u001b[39;00m\u001b[38;5;21;01mconfig\u001b[39;00m\n\u001b[0;32m     85\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m:\n\u001b[1;32m---> 86\u001b[0m     \u001b[38;5;28;01mraise\u001b[39;00m \u001b[38;5;167;01mImportError\u001b[39;00m(\n\u001b[0;32m     87\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCould not import chromadb python package. \u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m     88\u001b[0m         \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mPlease install it with `pip install chromadb`.\u001b[39m\u001b[38;5;124m\"\u001b[39m\n\u001b[0;32m     89\u001b[0m     )\n\u001b[0;32m     91\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m client \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[0;32m     92\u001b[0m     \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_client_settings \u001b[38;5;241m=\u001b[39m client_settings\n",
      "\u001b[1;31mImportError\u001b[0m: Could not import chromadb python package. Please install it with `pip install chromadb`."
     ]
    }
   ],
   "source": [
    "from langchain.chains import LLMChain,StuffDocumentsChain\n",
    "from langchain.document_transformers import (\n",
    "    LongContextReorder\n",
    ")\n",
    "from langchain.embeddings import HuggingFaceBgeEmbeddings\n",
    "from langchain.vectorstores import Chroma\n",
    "\n",
    "#使用huggingface托管的开源LLM来做嵌入，MiniLM-L6-v2是一个较小的LLM \n",
    "embedings = HuggingFaceBgeEmbeddings(model_name=\"all-MiniLM-L6-v2\")\n",
    "\n",
    "text = [\n",
    "    \"篮球是一项伟大的运动。\",\n",
    "    \"带我飞往月球是我最喜欢的歌曲之一。\",\n",
    "    \"凯尔特人队是我最喜欢的球队。\",\n",
    "    \"这是一篇关于波士顿凯尔特人的文件。\",\n",
    "    \"我非常喜欢去看电影。\",\n",
    "    \"波士顿凯尔特人队以20分的优势赢得了比赛。\",\n",
    "    \"这只是一段随机的文字。\",\n",
    "    \"《艾尔登之环》是过去15年最好的游戏之一。\",\n",
    "    \"L.科内特是凯尔特人队最好的球员之一。\",\n",
    "    \"拉里.伯德是一位标志性的NBA球员\"\n",
    "]\n",
    "\n",
    "retrieval = Chroma.from_texts(text,embedings).as_retriever(\n",
    "    search_kwargs={\"k\": 10}\n",
    ")\n",
    "query = \"关于凯尔特人队你都知道什么?\"\n",
    "\n",
    "#根据相关性返回文本块\n",
    "docs = retrieval.get_relevant_documents(query)\n",
    "docs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0ab4ef4-3a73-44cd-b327-9ff50f5bfef2",
   "metadata": {},
   "source": [
    "## 对检索结果进行重排序\n",
    "根据论文结果，当相关信息出现在输入上下文头尾相关性最高，中间相关性最低\n",
    "- 问题相关性越低的内容块放在中间\n",
    "- 问题相关性越高的内容块放在头尾"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ed71fe51-2a94-43cb-9694-b09d5b7c4b69",
   "metadata": {},
   "outputs": [],
   "source": [
    "reordering = LongContextReorder()\n",
    "reo_docs = reordering.transform_documents(docs) # 这个函数就是将相关性高的放在头尾\n",
    "reo_docs"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d5f7c858-4c8f-41ff-89ab-301ac9704b47",
   "metadata": {},
   "outputs": [],
   "source": [
    "#检测下这种方案的精度效果\n",
    "\n",
    "from dotenv import load_dotenv\n",
    "load_dotenv(\"openai.env\")\n",
    "import os\n",
    "api_key = os.environ.get(\"OPENAI_API_KEY\")\n",
    "api_base = os.environ.get(\"OPENAI_API_BASE\")\n",
    "\n",
    "\n",
    "from langchain.prompts import PromptTemplate\n",
    "from langchain.llms import OpenAI\n",
    "\n",
    "#设置llm\n",
    "llm = OpenAI(\n",
    "    api_key=api_key,\n",
    "    api_base=api_base,\n",
    "    model=\"gpt-3.5-turbo-instruct\",\n",
    "    temperature=0\n",
    ")\n",
    "\n",
    "document_prompt = PromptTemplate(\n",
    "    input_variables=[\"page_content\"],template=\"{page_content}\"\n",
    ")\n",
    "\n",
    "stuff_prompt_override =\"\"\"Given this text extracts:\n",
    "----------------------------------------\n",
    "{context}\n",
    "----------------------------------------\n",
    "Please answer the following questions:\n",
    "{query}\n",
    "\"\"\"\n",
    "\n",
    "prompt = PromptTemplate(\n",
    "    template=stuff_prompt_override,\n",
    "    input_variables=[\"context\",\"query\"]\n",
    ")\n",
    "\n",
    "llm_chain = LLMChain(\n",
    "    llm=llm,\n",
    "    prompt=prompt\n",
    ")\n",
    "\n",
    "WorkChain = StuffDocumentsChain(\n",
    "    llm_chain=llm_chain,\n",
    "    document_prompt=document_prompt,\n",
    "    document_variable_name=\"context\"\n",
    ")\n",
    "\n",
    "#调用\n",
    "WorkChain.run(\n",
    "    input_documents=reo_docs,\n",
    "    query=\"我最喜欢做什么事情？\"\n",
    ")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python (langchain)",
   "language": "python",
   "name": "langchain-env"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.13.0"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
